# -*- coding: utf-8 -*-
"""
断点管理器
负责保存和恢复流水线执行状态
"""

import json
import logging
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, Optional

from .step_base import StepStatus

logger = logging.getLogger(__name__)


class CheckpointManager:
    """断点管理器"""

    def __init__(self, checkpoint_dir: Path):
        """
        初始化断点管理器

        Args:
            checkpoint_dir: 断点目录
        """
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)

        # 断点文件路径
        self.pipeline_file = self.checkpoint_dir / "pipeline_state.json"
        self.steps_dir = self.checkpoint_dir / "steps"

        # 创建步骤目录
        self.steps_dir.mkdir(exist_ok=True)

    def save_pipeline_state(
        self,
        pipeline_name: str,
        current_step: Optional[str],
        steps_state: Dict[str, Any],
        metadata: Optional[Dict[str, Any]] = None,
    ):
        """
        保存流水线状态

        Args:
            pipeline_name: 流水线名称
            current_step: 当前步骤名
            steps_state: 步骤状态
            metadata: 元数据
        """
        state = {
            "pipeline_name": pipeline_name,
            "current_step": current_step,
            "timestamp": datetime.now().isoformat(),
            "steps": steps_state,
            "metadata": metadata or {},
        }

        with open(self.pipeline_file, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2, ensure_ascii=False)

        logger.info(f"Pipeline state saved: {pipeline_name}")

    def load_pipeline_state(self) -> Optional[Dict[str, Any]]:
        """
        加载流水线状态

        Returns:
            状态字典，如果无断点则返回None
        """
        if not self.pipeline_file.exists():
            logger.info("No pipeline state file found")
            return None

        try:
            with open(self.pipeline_file, "r", encoding="utf-8") as f:
                state = json.load(f)

            logger.info(f"Pipeline state loaded: {state['pipeline_name']}")
            return state

        except Exception as e:
            logger.error(f"Failed to load pipeline state: {e}")
            return None

    def save_step_state(
        self,
        step_name: str,
        step_status: StepStatus,
        result: Optional[Dict[str, Any]] = None,
    ):
        """
        保存步骤状态

        Args:
            step_name: 步骤名称
            step_status: 步骤状态
            result: 步骤结果
        """
        step_file = self.steps_dir / f"{step_name}.json"

        state = {
            "step_name": step_name,
            "status": step_status.value,
            "timestamp": datetime.now().isoformat(),
            "result": result.to_dict() if result else None,
        }

        with open(step_file, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2, ensure_ascii=False)

        logger.debug(f"Step state saved: {step_name} ({step_status.value})")

    def load_step_state(self, step_name: str) -> Optional[Dict[str, Any]]:
        """
        加载步骤状态

        Args:
            step_name: 步骤名称

        Returns:
            状态字典
        """
        step_file = self.steps_dir / f"{step_name}.json"

        if not step_file.exists():
            logger.debug(f"No step state file found: {step_name}")
            return None

        try:
            with open(step_file, "r", encoding="utf-8") as f:
                state = json.load(f)

            logger.debug(f"Step state loaded: {step_name}")
            return state

        except Exception as e:
            logger.error(f"Failed to load step state: {step_name} - {e}")
            return None

    def clear_checkpoint(self, pipeline_name: Optional[str] = None):
        """
        清理断点

        Args:
            pipeline_name: 流水线名称（用于验证）
        """
        if pipeline_name:
            state = self.load_pipeline_state()
            if state and state.get("pipeline_name") != pipeline_name:
                logger.warning(f"Pipeline name mismatch: {pipeline_name} vs {state.get('pipeline_name')}")
                return

        # 删除断点文件
        for file in [self.pipeline_file] + list(self.steps_dir.glob("*.json")):
            if file.exists():
                file.unlink()

        logger.info("Checkpoint cleared")

    def list_checkpoints(self) -> list[Dict[str, Any]]:
        """
        列出所有断点

        Returns:
            断点列表
        """
        checkpoints = []

        if not self.pipeline_file.exists():
            return checkpoints

        try:
            state = self.load_pipeline_state()
            if state:
                checkpoints.append({
                    "pipeline_name": state["pipeline_name"],
                    "current_step": state["current_step"],
                    "timestamp": state["timestamp"],
                    "steps_count": len(state.get("steps", {})),
                })
        except Exception as e:
            logger.error(f"Failed to list checkpoints: {e}")

        return checkpoints

    def has_checkpoint(self) -> bool:
        """
        检查是否有断点

        Returns:
            是否有断点
        """
        return self.pipeline_file.exists()
